from typing import List, Union, Optional

import tiktoken  # OpenAI 开发的一个用于计算文本 token 数量的 Python 库
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage


class TokenCounter:
    def __init__(self, model):
        _default_tokenizer = tiktoken.get_encoding("cl100k_base")
        if model is None:
            self.tokenizer = _default_tokenizer
        else:
            try:
                self.tokenizer = tiktoken.encoding_for_model(model.model_name)
            except (KeyError, TypeError):
                self.tokenizer = _default_tokenizer

    def count_text(self, text: str) -> int:
        """Calculate tokens for a text string"""
        return 0 if not text else len(self.tokenizer.encode(text))

    def count_msg(self, msg: List[Union[dict, HumanMessage]]) -> int:
        """Calculate tokens for a list of messages"""
        return sum(self.count_text(item['content'] if isinstance(item, dict) else item.content) for item in msg)

    def count_one_round(self, input_msg: List[Union[dict, HumanMessage, SystemMessage]], output_msg: AIMessage) -> \
            tuple[Optional[int], Optional[int]]:

        input_tokens: Optional[int] = None
        output_tokens: Optional[int] = None

        if hasattr(output_msg, "usage_metadata") and output_msg.usage_metadata:
            # 优先取AI返回的实际token数
            input_tokens = output_msg.usage_metadata.get("input_tokens", None)
            output_tokens = output_msg.usage_metadata.get("output_tokens", None)

        if input_tokens is None:
            # 如果AIMessage未提供input_tokens，则通过文本计算
            input_tokens = self.count_msg(input_msg)

        if output_tokens is None:
            # 如果AIMessage未提供output_tokens，则通过文本计算
            output_tokens = self.count_text(output_msg.content)
            output_tokens += self.count_text(output_msg.additional_kwargs.get('reasoning_content', ''))

        return input_tokens, output_tokens
